import os
import xml.etree.ElementTree as ET
from pathlib import Path

class YoloVocConverter:
    def __init__(self):
        self.classes = []
        self.class_to_id = {}
        
    def set_classes(self, classes):
        """设置类别列表并验证"""
        if not classes:
            raise ValueError("类别列表不能为空")
        if len(classes) != len(set(classes)):
            raise ValueError("类别名称不能重复")
            
        self.classes = classes
        self.class_to_id = {name: i for i, name in enumerate(classes)}
        
    def yolo_to_voc(self, yolo_txt, img_width, img_height):
        """将YOLO格式转换为VOC格式"""
        results = []
        for line in yolo_txt.splitlines():
            if not line.strip():
                continue
            class_id, x_center, y_center, width, height = map(float, line.split())
            
            # 验证类别ID
            if int(class_id) >= len(self.classes):
                raise ValueError(f"类别ID {int(class_id)} 超出范围（最大 {len(self.classes)-1}）")
            
            # 计算边界框坐标
            x_min = (x_center - width / 2) * img_width
            y_min = (y_center - height / 2) * img_height
            x_max = (x_center + width / 2) * img_width
            y_max = (y_center + height / 2) * img_height
            
            results.append({
                'class': self.classes[int(class_id)],
                'xmin': int(x_min),
                'ymin': int(y_min),
                'xmax': int(x_max),
                'ymax': int(y_max)
            })
        return results
        
    def voc_to_yolo(self, voc_annots, img_width, img_height):
        """将VOC格式转换为YOLO格式"""
        results = []
        for annot in voc_annots:
            class_name = annot.find('name').text
            if class_name not in self.class_to_id:
                raise ValueError(f"未知类别: {class_name}")
                
            class_id = self.class_to_id[class_name]
            
            bbox = annot.find('bndbox')
            xmin = float(bbox.find('xmin').text)
            ymin = float(bbox.find('ymin').text)
            xmax = float(bbox.find('xmax').text)
            ymax = float(bbox.find('ymax').text)
            
            # 计算YOLO格式坐标
            x_center = ((xmin + xmax) / 2) / img_width
            y_center = ((ymin + ymax) / 2) / img_height
            width = (xmax - xmin) / img_width
            height = (ymax - ymin) / img_height
            
            results.append(f"{class_id} {x_center} {y_center} {width} {height}")
        return results
        
    def parse_voc_xml(self, xml_path):
        """解析VOC XML文件"""
        tree = ET.parse(xml_path)
        root = tree.getroot()
        return root.findall('object')
        
    def save_voc_xml(self, img_path, size, objects, output_path):
        """保存为VOC XML格式"""
        root = ET.Element('annotation')
        
        # 添加文件信息
        folder = ET.SubElement(root, 'folder')
        folder.text = str(Path(img_path).parent.name)
        
        filename = ET.SubElement(root, 'filename')
        filename.text = str(Path(img_path).name)
        
        # 添加图像尺寸
        size_elem = ET.SubElement(root, 'size')
        width = ET.SubElement(size_elem, 'width')
        width.text = str(size[0])
        height = ET.SubElement(size_elem, 'height')
        height.text = str(size[1])
        depth = ET.SubElement(size_elem, 'depth')
        depth.text = '3'
        
        # 添加每个对象
        for obj in objects:
            obj_elem = ET.SubElement(root, 'object')
            name = ET.SubElement(obj_elem, 'name')
            name.text = obj['class']
            
            bndbox = ET.SubElement(obj_elem, 'bndbox')
            xmin = ET.SubElement(bndbox, 'xmin')
            xmin.text = str(obj['xmin'])
            ymin = ET.SubElement(bndbox, 'ymin')
            ymin.text = str(obj['ymin'])
            xmax = ET.SubElement(bndbox, 'xmax')
            xmax.text = str(obj['xmax'])
            ymax = ET.SubElement(bndbox, 'ymax')
            ymax.text = str(obj['ymax'])
            
        # 保存XML文件
        tree = ET.ElementTree(root)
        tree.write(output_path, encoding='utf-8', xml_declaration=True)
